library(Rcpp)
library(dplyr)
library(tidyr) 
library(abind)

set.seed(123)
rm(list=ls())

### load the functions needed to compute negative log likelihood
Rcpp::sourceCpp("src/misc.cpp")

### set subfolders
in_dir <- "./data"
aux_dir <- "./misc_inputs"
outdir <- "./output"

### load data and estimates
(loaded <-  load(file.path(in_dir, "analysis_data.RData"))) 
est <- readRDS(file.path(outdir, "rum_estimates.rds"))

### set constants
countries <- c("CZ","DE","ES","FR","HU","IT","NL","SE","GB")
names(countries) <- countries

countries.lab <- countrycode::countrycode(countries, 'iso2c', 'country.name')
names(countries.lab) <- countries
countries.lab['GB'] <- 'England'
countries.lab <- sort(countries.lab)

ucov <- c("female", "university", "age_under30", "age_over50")
dims <- c("econ","cult","populism","ant","ppl","man")
dims.lab <- c("Economic Left-Right","Cultural Conservatism","Populism","Anti-elitism","People-centrism","Manichean outlook")
mdims <- c('econ', 'cult')
sdims <- c('econ', 'cult', 'populism')
fdims <- c("econ","cult", "ant","ppl","man")
popvars <- c('populism','ant','ppl','man')

### template for tables
t0 <- read.table(header=TRUE, sep = "|", text="
sorter|term|label
5|alpha|$\\text{logit}(\\delta)$
10|lambda1|$\\lambda_{1,1}$
15|lambda2|$\\lambda_{2,2}$
20|lambda3|$\\lambda_{3,3}$
25|lambda4|$\\lambda_{2,1}$
30|lambda5|$\\lambda_{3,2}$
35|lambda6|$\\lambda_{3,1}$
1000|AIC|AIC
1001|BIC|BIC
1002|nobs|N")

### inventorization
inventory <- data.frame(nam=names(est)) |>
  mutate( 
  cntry = substring(nam, 4L),
  model = substr(nam, 1L, 2L),
  model_no = substr(nam, 2L, 2L),
  conv = sapply(nam, function(u) est[[u]]@details$convergence),
  ndim = ifelse(substring(model,1L,1L)=='A', "3D", "2D") 
  ) |>
  filter(conv==0)
 
table(inventory$cntry,inventory$model) 

### model fit statistics
inventory$nobs <- inventory$BIC <- inventory$ll <- inventory$AIC <- NA
for (j in seq_along(inventory$nam)) {
  inventory[j,'nobs'] <- stats4::nobs(est[[inventory$nam[j]]])  
  inventory[j,'ll'] <- stats4::logLik(est[[inventory$nam[j]]])  
  inventory[j,'AIC'] <- stats4::AIC(est[[inventory$nam[j]]])  
  inventory[j,'BIC'] <- stats4::BIC(est[[inventory$nam[j]]])  
}

inventory |> select(model, cntry, ndim, nobs, ll, AIC, BIC) |>
  write.csv(file.path(outdir, "vc_model_fit.csv"), row.names=FALSE)

### find the covariates used in estimation
saml <- lapply(countries, function(cn) {
  v <- vdata[[cn]] |> tibble::rowid_to_column("rowid")
  p <- subset(pdata, cntry==cn)
  menu <- v |> count(cntry_party_ivc) |> filter(n>10) |> 
    pull(cntry_party_ivc) |> intersect(p$cntry_party) |> sort() 
  v <- subset(v, cntry_party_ivc %in% menu, select=c('rowid','wt','cntry_party_ivc', dims, ucov))
  
  # party-level intercepts and control variables
  tobin <- setdiff(menu, v |> count(cntry_party_ivc) |> slice_max(n) |> pull(cntry_party_ivc)) 
  xlab1 <- paste0(make.names(tobin), ": constant")
  for (uv in ucov) {
    xlab1 <- c(xlab1, paste0(make.names(tobin),": ",uv))
  }
  return(xlab1) 
})
 
## export summaries of the parameter estimates
fest <- lapply(countries, function(cn) {
  sel <- inventory |>  filter(cntry == cn)
  e <- lapply(seq_len(nrow(sel)), function(r) {
    m <- sel$nam[r]
    mod <- paste0("c", sel$model_no[r], '.', tolower(sel$ndim[r]))
    tab <- as.data.frame(stats4::summary(est[[m]])@coef)
    colnames(tab) <- c("val","set") 
    main <- tab |> mutate(term = rownames(tab)) |>
      mutate(
        term = rownames(tab),
        p = pnorm(-abs(val)/set),
        stars = case_when(p <0.01 ~ "***",  p <0.05 ~ "**",  p <0.1 ~ "*", TRUE ~ ""),
        val1 = paste0(format(round(val, digits=2), trim=TRUE), stars),
        val2 = paste0("(", format(round(set, digits=2), trim=TRUE), ")")
      ) |> select(term, val1, val2) |>
      pivot_longer(cols=starts_with("val"), names_to="subrow", values_to=mod)
    modsta <- data.frame(term = c("AIC","BIC","nobs"),
                         subrow = 'val1',
                         val = c(
                           format(round(sel$AIC[r], digits=1), trim=TRUE),
                           format(round(sel$BIC[r], digits=1), trim=TRUE),
                           format(round(sel$nobs[r], digits=0), trim=TRUE)
                         ))
    colnames(modsta)[3L] <- mod
    rbind(main, modsta)
  }) 
  t <- Reduce(function(x,y) merge(x,y, all=TRUE), e)
  t <- replace(t, is.na(t), "") |> left_join(t0)
  left <- saml[[cn]]
  names(left) <- paste0('beta', seq_along(left))
  sorter2 <- numeric(nrow(t))
  sorter2[grep("^beta", t$term)] <- as.numeric(substring(t$term[grep("^beta", t$term)], 5L))
  t <- t |> mutate(
    label = coalesce(label, left[term]),
    sorter = coalesce(sorter, 200 + sorter2),
    sorter = sorter + as.numeric(substring(subrow,4L))/10) |>
    arrange(sorter) |>
    mutate(
      label = ifelse(subrow=="val2", "", label),
      end = ifelse(subrow=="val1", "\\\\*", "\\\\")
    )
  cnams <- sort(grep("^c", colnames(t),val=TRUE))
  list(rows = paste0(apply(t[,c("label", cnams)], 1L, paste, collapse="&"), t$end), 
       header = paste0("Parameter&", paste(cnams, collapse="&"), "\\\\\\hline"),
       ncol = length(cnams))
})

write("\n", file = file.path(outdir, "est_rums.txt"))
for (cn in countries) {
  write(
    paste0("\n\n\\begin{longtable}{l", paste(rep("c", fest[[cn]]$ncol), collapse=""), "} \\caption{Covariates of vote choice, ",
           countries.lab[cn], "}\\\\
\\hline ", fest[[cn]]$header, "
\\endfirsthead
\\hline ", fest[[cn]]$header, "
\\endhead
\\hline\\multicolumn{",(fest[[cn]]$ncol+1L),"}{r}{Continued on next page}\\\\
\\endfoot
\\hline\\hline
\\hline \\multicolumn{",(fest[[cn]]$ncol+1L),"}{r}{*p<0.1, **p<0.05, ***p<0.01. Standard errors in parentheses.}\\\\
\\endlastfoot
"), file = file.path(outdir, "est_rums.txt"), append=TRUE)
  write(fest[[cn]]$rows, file = file.path(outdir, "est_rums.txt"), append=TRUE)
  write("\\end{longtable}", file = file.path(outdir, "est_rums.txt"), append=TRUE)
}
